import torch
from torch import nn
import math


class ChannelAttention(nn.Module):
    def __init__(self, channels, ratio=16, gamma=2, b=1):
        super(ChannelAttention, self).__init__()
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.max_pool = nn.AdaptiveAvgPool2d(1)

        # 输入通道数自适应修改卷积核大小
        kernel_size = int(abs((math.log(channels, 2) + b) / gamma))
        kernel_size = kernel_size if kernel_size % 2 else kernel_size + 1
        padding = kernel_size // 2
        self.avg_pool = nn.AdaptiveAvgPool2d(1)
        self.conv = nn.Conv1d(1, 1, kernel_size, padding=padding, bias=False)

        self.fc = nn.Sequential(
            nn.Linear(channels, channels // ratio, bias=False),
            nn.ReLU(),
            nn.Linear(channels // ratio, channels, bias=False),
        )
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        b, c, h, w = x.size()
        max_pool_out = self.avg_pool(x).view(b, 1, c)
        avg_pool_out = self.max_pool(x).view(b, 1, c)
        avg = self.avg_pool(x).view([b, 1, c])
        conv_out1 = self.conv(max_pool_out)
        conv_out2 = self.conv(avg_pool_out)
        out = conv_out1 + conv_out2
        out = self.sigmoid(out).view([b, c, 1, 1])
        return out * x


class SpaceAttention(nn.Module):
    def __init__(self, kernel_size=7):
        super(SpaceAttention, self).__init__()
        padding = kernel_size // 2
        self.conv = nn.ConvTranspose2d(2, 1, kernel_size, 1, padding, bias=False)
        self.sigmoid = nn.Sigmoid()

    def forward(self, x):
        max_pool_out, _ = torch.max(x, dim=1, keepdim=True)
        mean_pool_out = torch.mean(x, dim=1, keepdim=True)
        pool_out = torch.cat([max_pool_out, mean_pool_out], dim=1)
        out = self.conv(pool_out)
        out = self.sigmoid(out)
        return out * x


class CbamECA(nn.Module):
    def __init__(self, channel, ratio=16, kernel_size=7):
        super(CbamECA, self).__init__()
        self.channel_attention = ChannelAttention(channel, ratio)
        self.space_attention = SpaceAttention(kernel_size)

    def forward(self, x):
        x = self.channel_attention(x)
        x = self.space_attention(x)
        return x
